//
// Created by wangwanlong on 2024/6/24.
//

#ifndef MCGPP_PARTITION_H
#define MCGPP_PARTITION_H

#include <array>
#include "graph.h"
#include "graph_io.h"
#include "partition/coarsen/ml_coarsener.h"
#include "context.h"
#include "sa_refiner.h"
#include "tabu_refiner.h"
#include "desent_refiner.h"
#include "policies/stop_policy.h"
#include "metrics.h"
#include "coarsen/mem_coarsener.h"
#include "timer.h"

namespace mcgpp {
  class Partitioner {
  public:
    Partitioner(ds::GenericGraph &graph) : _g(graph) {
    }

    void partition() {
      //读取种子分区
      readPartitionFile();
      double c = metrics::conductance(_g);
      LOG << "before coarsen init conductance:" << c;

      //缩合
      MemCoarnsener<> memCoarsener(_g);
      memCoarsener.coarsenImpl(Context::instance().coarsen_limit);

      c = metrics::conductance(_g);
      LOG << "after coarsen init conductance:" << c;

      //局部搜索与还原
      SaRefiner saRefiner(_g);
      saRefiner.initializeImpl();
      saRefiner.refineImpl();

      memCoarsener.uncoarsenImpl();


    }

    void saveGlobalSol() {

    }

    void vcyclePartition() {
      double best_time;//找到最优解的时间
      Vec<IDType> result(_g.initialNumNodes());
      ll best_cond;
      ll current_cond;
      //读取种子分区
      readPartitionFile();
      //记录初始电导
      Context::instance().partitionResult.init_conductance = metrics::conductance(_g);
      current_cond = Context::instance().partitionResult.init_conductance * 1e10;
      MemCoarnsener<> memCoarsener(_g);
      SaRefiner saRefiner(_g);
      TabuRefiner tabuRefiner(_g);
      int level = 0;
      //保存历史最优结果
      auto  saveBestResul = [&](ll cond){
        best_time = Timer::instance().getCostTime();
        best_cond = cond;
        for (auto v: _g._current_node_set)
          result[v] = _g.partID(v);
      };
      //保存初始结果
      saveBestResul(Context::instance().partitionResult.init_conductance * 1e10);
      while (!Timer::instance().TLE()) {
        LOG << "*****************************************************************";
        LOG << "begin vcycle iter" << level;
        if (Context::instance().coarsen_limit >= _g.initialNumNodes()) {
          saRefiner.initializeImpl();
          saRefiner.refineImpl();
          tabuRefiner.initializeImpl();
          tabuRefiner.refineImpl();
        } else {
          LOG<<"---coarsening-----";
          memCoarsener.coarsenImpl(Context::instance().coarsen_limit, saRefiner,tabuRefiner);
          LOG<<"---uncoarsening-----";
          memCoarsener.uncoarsenImpl(saRefiner,tabuRefiner);
        }
        //保存历史最优解
        if (tabuRefiner.current_obj.conductance < best_cond) {
          saveBestResul(tabuRefiner.current_obj.conductance);
        }

//        if(Randomize::instance().getRandomInt(0,1) == 0){
//          //切换
//          if(Context::instance().objMetric == ObjMetric::cut_weight)
//            Context::instance().objMetric = ObjMetric::cut_num;
//          else
//            Context::instance().objMetric = ObjMetric::cut_weight;
//        }

        LOG << "best cond = " << std::fixed << std::setprecision(10)<< static_cast<double>(tabuRefiner.current_obj.conductance / 1e10);
        LOG << "opt cond = " << std::fixed << std::setprecision(10) << static_cast<double>(best_cond / 1e10);

        ++level;
      }
      LOG << "total level = " << level;
      Context::instance().partitionResult.best_time = best_time;
      Context::instance().partitionResult.total_time = Timer::instance().getCostTime();
      Context::instance().partitionResult.conductance = best_cond / 1e10;
      Context::instance().partitionResult.result = std::move(result);

      //验证结果
      verifySol(best_cond, Context::instance().partitionResult.result);
    }


    //没有初始解的V循环，看看梯度下降算法的作用
    void vcyclePartition2(){
      double best_time;//找到最优解的时间
      Vec<IDType> result(_g.initialNumNodes());
      Metrics best_obj;

      MemCoarnsener<> memCoarsener(_g);
      SaRefiner saRefiner(_g);
      int level = 0;
      //保存历史最优结果
      auto  saveBestResul = [&](){
        best_time = Timer::instance().getCostTime();
        best_obj = saRefiner.current_obj;
        for (auto v: _g._current_node_set)
          result[v] = _g.partID(v);
      };
      saveBestResul();

      MLCoarsener<>mlCoarsener(_g);//只缩合

      while (!Timer::instance().TLE()){
        LOG<<"---coarsening-----";
        mlCoarsener.coarsenImpl(5000);




      }
    }

    //n-level缩合，还原一对顶点后搜索
    void vcyclePartition3() {
      double best_time;//找到最优解的时间
      Vec<IDType> result(_g.initialNumNodes());
      ll best_cond;
      ll current_cond;
      //读取种子分区
      readPartitionFile();
      //记录初始电导
      Context::instance().partitionResult.init_conductance = metrics::conductance(_g);
      current_cond = Context::instance().partitionResult.init_conductance * 1e10;

      MLCoarsener<>mlCoarsener(_g);

      DesentRefiner<AdvancedRandomWalkModelStopsSearch> desentRefiner(_g);

      int level = 0;
      //保存历史最优结果
      auto  saveBestResule = [&](ll cond){
        best_time = Timer::instance().getCostTime();
        best_cond = cond;
        for (auto v: _g._current_node_set)
          result[v] = _g.partID(v);
      };
      Metrics f_best;
      //保存初始结果
      saveBestResule(Context::instance().partitionResult.init_conductance * 1e10);
      while (!Timer::instance().TLE()) {
        LOG << "*****************************************************************";
        LOG << "begin vcycle iter" << level;
        LOG<<"---coarsening-----";
        mlCoarsener.coarsenImpl(Context::instance().coarsen_limit);
        LOG<<"---uncoarsening-----";
        mlCoarsener.uncoarsenImpl(desentRefiner,f_best);

        if(f_best.conductance < best_cond){
          saveBestResule(f_best.conductance);
        }


        LOG << "best cond = " << std::fixed << std::setprecision(10)<< static_cast<double>(f_best.conductance / 1e10);
        LOG << "opt cond = " << std::fixed << std::setprecision(10) << static_cast<double>(best_cond / 1e10);

        ++level;
      }
      LOG << "total level = " << level;
      Context::instance().partitionResult.best_time = best_time;
      Context::instance().partitionResult.total_time = Timer::instance().getCostTime();
      Context::instance().partitionResult.conductance = best_cond / 1e10;
      Context::instance().partitionResult.result = std::move(result);

      //验证结果
      verifySol(best_cond, Context::instance().partitionResult.result);
    }

    //不使用混合搜索
    void vcyclePartition4() {
      double best_time;//找到最优解的时间
      Vec<IDType> result(_g.initialNumNodes());
      ll best_cond;
      ll current_cond;
      //读取种子分区
      readPartitionFile();
      //记录初始电导
      Context::instance().partitionResult.init_conductance = metrics::conductance(_g);
      current_cond = Context::instance().partitionResult.init_conductance * 1e10;
      MemCoarnsener<> memCoarsener(_g);
      SaRefiner saRefiner(_g);
      TabuRefiner tabuRefiner(_g);
      int level = 0;
      //保存历史最优结果
      auto  saveBestResul = [&](ll cond){
        best_time = Timer::instance().getCostTime();
        best_cond = cond;
        for (auto v: _g._current_node_set)
          result[v] = _g.partID(v);
      };
      //保存初始结果
      saveBestResul(Context::instance().partitionResult.init_conductance * 1e10);
      while (!Timer::instance().TLE()) {
        LOG << "*****************************************************************";
        LOG << "begin vcycle iter" << level;
        if (Context::instance().coarsen_limit >= _g.initialNumNodes()) {
          saRefiner.initializeImpl();
          saRefiner.refineImpl();
          tabuRefiner.initializeImpl();
          tabuRefiner.refineImpl();
        } else {
          LOG<<"---coarsening-----";
          memCoarsener.coarsenImpl(Context::instance().coarsen_limit,tabuRefiner);
          LOG<<"---uncoarsening-----";
          memCoarsener.uncoarsenImpl(tabuRefiner);
        }
        //保存历史最优解
        if (tabuRefiner.current_obj.conductance < best_cond) {
          saveBestResul(tabuRefiner.current_obj.conductance);
        }

//        if(Randomize::instance().getRandomInt(0,1) == 0){
//          //切换
//          if(Context::instance().objMetric == ObjMetric::cut_weight)
//            Context::instance().objMetric = ObjMetric::cut_num;
//          else
//            Context::instance().objMetric = ObjMetric::cut_weight;
//        }

        LOG << "best cond = " << std::fixed << std::setprecision(10)<< static_cast<double>(tabuRefiner.current_obj.conductance / 1e10);
        LOG << "opt cond = " << std::fixed << std::setprecision(10) << static_cast<double>(best_cond / 1e10);

        ++level;
      }
      LOG << "total level = " << level;
      Context::instance().partitionResult.best_time = best_time;
      Context::instance().partitionResult.total_time = Timer::instance().getCostTime();
      Context::instance().partitionResult.conductance = best_cond / 1e10;
      Context::instance().partitionResult.result = std::move(result);

      //验证结果
      verifySol(best_cond, Context::instance().partitionResult.result);
    }

    void vcyclePartitionCheck() {

      double best_time;//找到最优解的时间
      Vec<IDType> result(_g.initialNumNodes());
      Metrics best_obj;

      //读取种子分区
      readPartitionFile();
      MemCoarnsener<> memCoarsener(_g);
      SaRefiner saRefiner(_g);
      int level = 0;
      _g.check();
      while (!Timer::instance().TLE()) {
        LOG << "*****************************************************************";
        LOG << "begin vcycle iter" << level;
        if (Context::instance().coarsen_limit >= _g.initialNumNodes()) {
          saRefiner.initializeImpl();
          saRefiner.refineImpl();
        } else {
          memCoarsener.coarsenImpl(Context::instance().coarsen_limit);
          memCoarsener.uncoarsenImpl();
        }

        //保存历史最优解
        if (saRefiner.current_obj.conductance < best_obj.conductance) {
          best_time = Timer::instance().getCostTime();
          best_obj = saRefiner.current_obj;
          if (_g._current_node_set.size() != _g.initialNumNodes())
            ERROREXIT("error!");
          for (auto v: _g._current_node_set)
            result[v] = _g.partID(v);

        }

        LOG << "best cond = " << static_cast<double>(saRefiner.current_obj.conductance / 1e10);
        LOG << "opt cond = " << static_cast<double>(best_obj.conductance / 1e10);

        ++level;
      }
      LOG << "total level = " << level;
      Context::instance().partitionResult.best_time = best_time;
      Context::instance().partitionResult.total_time = Timer::instance().getCostTime();
      Context::instance().partitionResult.conductance = best_obj.conductance / 1e10;
      Context::instance().partitionResult.result = std::move(result);

      //验证结果
      verifySol(best_obj.conductance, Context::instance().partitionResult.result);
    }

    //设置种子分区
    void readPartitionFile() {
      std::vector<IDType> part(_g.initialNumNodes());
      io::readPartitionFile(Context::instance().seed_partition_file, part);
      for (int i = 0; i < _g.initialNumNodes(); ++i) {
        _g.setNodePart(i, part[i]);
      }

    }

    //验证结果
    void verifySol(ll &best_cond, Vec<IDType> &result) {
      int node_num, edge_num;
      Vec<Vec<IDType>> adj_list;
      io::readGraph(Context::instance().instance_file, node_num, edge_num, adj_list);

      std::array<ll, 2> vol = {0, 0};
      ll cut = 0;
      int sum = 0;
      for (int i = 0; i < node_num; ++i) {
        for (auto adj: adj_list[i]) {
          if (result[i] != result[adj])
            ++cut;
        }
        vol[result[i]] += adj_list[i].size();
        sum += adj_list[i].size();
      }
      if (vol[0] + vol[1] != edge_num * 2) {
        LOG << sum << (vol[0] + vol[1]);
        ERROREXIT("vol error!");
      }
      cut /= 2;
      ll cond = static_cast<double>((double) cut / (double) std::min(vol[0], vol[1])) * 1e10;

      if (cond != best_cond) {
        LOG << cond << cut << vol[0] << vol[1];
        LOG << best_cond;
        ERROREXIT("cond error!");
      }

      double tmp = best_cond / 1e10;
      LOG << "check result success!" << "final result:" << std::fixed << std::setprecision(10) << tmp;


    }


    ds::GenericGraph &_g;
  };
}
#endif //MCGPP_PARTITION_H
